"""
数据库工具模块
"""
import time
from loguru import logger
import pymysql
from sqlalchemy import text

from config.database import MYSQL_CONFIG
from database.db_engine import engine, init_db

def create_database_if_not_exists():
    """
    检查并创建数据库（如果不存在）
    
    Returns:
        bool: 操作是否成功
    """
    try:
        # 不指定数据库名连接MySQL
        conn = pymysql.connect(
            host=MYSQL_CONFIG['host'],
            port=MYSQL_CONFIG['port'],
            user=MYSQL_CONFIG['user'],
            password=MYSQL_CONFIG['password'],
            charset=MYSQL_CONFIG['charset']
        )
        
        # 创建游标
        cursor = conn.cursor()
        
        # 检查数据库是否存在
        cursor.execute(f"SHOW DATABASES LIKE '{MYSQL_CONFIG['database']}'")
        result = cursor.fetchone()
        
        # 如果数据库不存在，则创建
        if not result:
            logger.info(f"数据库 {MYSQL_CONFIG['database']} 不存在，正在创建...")
            cursor.execute(f"CREATE DATABASE {MYSQL_CONFIG['database']} CHARACTER SET {MYSQL_CONFIG['charset']}")
            logger.info(f"数据库 {MYSQL_CONFIG['database']} 创建成功")
        else:
            logger.info(f"数据库 {MYSQL_CONFIG['database']} 已存在")
        
        # 关闭连接
        cursor.close()
        conn.close()
        return True
        
    except Exception as e:
        logger.error(f"创建数据库失败: {e}")
        return False

def initialize_database():
    """
    初始化数据库，创建表结构
    
    Returns:
        bool: 操作是否成功
    """
    try:
        # 先确保数据库存在
        if not create_database_if_not_exists():
            return False
        
        # 创建所有表
        init_db()
        logger.info("数据库表结构初始化成功")
        return True
        
    except Exception as e:
        logger.error(f"初始化数据库失败: {e}")
        return False

def clear_old_realtime_data(days=45):
    """
    清理过旧的实时数据，避免数据库过大
    
    Args:
        days: 保留最近多少天的数据
        
    Returns:
        bool: 操作是否成功
    """
    try:
        # 执行删除SQL
        with engine.connect() as conn:
            sql = f"DELETE FROM stock_realtime_data WHERE DATEDIFF(NOW(), trade_time) > {days}"
            result = conn.execute(text(sql))
            deleted_count = result.rowcount
            conn.commit()
            
        logger.info(f"成功清理 {deleted_count} 条过期实时数据（{days}天前）")
        return True
        
    except Exception as e:
        logger.error(f"清理过期实时数据失败: {e}")
        return False

def clear_old_daily_data(days=365):
    """
    清理过旧的每日交易数据，避免数据库过大
    
    Args:
        days: 保留最近多少天的数据
        
    Returns:
        bool: 操作是否成功
    """
    try:
        # 执行删除SQL
        with engine.connect() as conn:
            sql = f"DELETE FROM stock_daily_data WHERE DATEDIFF(NOW(), trade_date) > {days}"
            result = conn.execute(text(sql))
            deleted_count = result.rowcount
            conn.commit()
            
        logger.info(f"成功清理 {deleted_count} 条过期每日交易数据（{days}天前）")
        return True
        
    except Exception as e:
        logger.error(f"清理过期每日交易数据失败: {e}")
        return False 
    
def clear_old_data(realtime_days=45, daily_days=365):
    """
    清理过旧的实时数据和每日交易数据，避免数据库过大
    Args:
        realtime_days: 保留最近多少天的实时数据
        daily_days: 保留最近多少天的每日交易数据
        
    Returns:
        bool: 操作是否成功
    """
    try:
        # 清理实时数据
        clear_old_realtime_data(realtime_days)
        # 清理每日交易数据
        clear_old_daily_data(daily_days)
        return True
    except Exception as e:
        logger.error(f"清理过旧数据失败: {e}")
        return False    